# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

from quantlib import *

def makeGenerator(maturity, dimension, seed, initial, vol, div, rate):
    uniformChoice = 'MersenneTwister'
    gaussianChoice = 'InverseNormal'
    stochasticChoice = 'BlackScholesMerton'

    brownianBridge = False

    # # Random number generator    
    if uniformChoice == 'MersenneTwister' and gaussianChoice == 'InverseNormal':
        randomGenerator = GetGenericMersenneTwisterInverseNormal(dimension,seed)
    else:
        raise Exception('Generator not supplied yet')

    # # Process
    if stochasticChoice == 'BlackScholesMerton':
        todaysDate = Date(15,May,1998)
        settlementDate = Date(17,May,1998)
        underlying = SimpleQuote(initial)
        volatility = BlackConstantVol(todaysDate, TARGET(), vol, Actual365Fixed())
        dividendYield = FlatForward(settlementDate, div, Actual365Fixed())
        riskFreeRate = FlatForward(settlementDate, rate, Actual365Fixed())

        x=QuoteHandle(underlying)
        x1=YieldTermStructureHandle(dividendYield)
        x2=YieldTermStructureHandle(riskFreeRate)
        x3=BlackVolTermStructureHandle(volatility)
        process = BlackScholesMertonProcess(x,x1,x2,x3)
    else:
        raise Exception('Process not supplied yet')

    # # Path
    pathGenerator = MersenneInverseNormalGenerator(process, maturity, dimension, randomGenerator, brownianBridge)
    return pathGenerator

##    return pathGenerator
##    sample = pathGenerator.next()
##    path = sample.value
##    print('First sample:')
##    for idx in range(dimension+1):
##        print(path.at(idx))
##    sample = pathGenerator.next()
##    path = sample.value
##    print('Second sample:')
##    for idx in range(dimension+1):
##        print(path.at(idx))



if __name__ == '__main__':
    maturity = 100.0;
    noOfPathsToPlot = 1 #set to 0 if no plotting wanted
    dimension = 100;
    seed = 15;
    initial = 100.0
    vol = 0.15
    div = 0.0
    rate = 0.05
    
    pathGenerator = makeGenerator(maturity, dimension, seed, initial, vol, div, rate)

    if noOfPathsToPlot > 0:
        import matplotlib.pyplot as plt
        fig = plt.figure()
        ax = fig.add_subplot(111)
        xvals = range(dimension+1)
        yvals = [0.0 for jdx in range(dimension+1)]
        
        for idx in range(noOfPathsToPlot):
            sample = pathGenerator.next()
            path = sample.value
            for jdx in range(dimension+1):
                yvals[jdx] = path.at(jdx)
        ax.plot(xvals, yvals)
        plt.show()
